Python: feat: add Amazon Bedrock Knowledge Base tool and context provider#7066
Draft
PVidyadhar wants to merge 1 commit into
Draft
Python: feat: add Amazon Bedrock Knowledge Base tool and context provider#7066PVidyadhar wants to merge 1 commit into
PVidyadhar wants to merge 1 commit into
Conversation
- Created BedrockKnowledgeBaseTool with async run() + get_tool_definition() - Created BedrockKnowledgeBaseProvider (ContextProvider subclass) with before_run() - Two integration points: standalone tool + automatic context injection - Supports managed search and agentic retrieval with fallback - Unit tests included - Added BEDROCK_MANAGED_KB.md design doc
giles17
added
documentation
github-actions
Bot
changed the title
|
@PVidyadhar please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds Amazon Bedrock Knowledge Base retrieval capabilities to the Python tools package, enabling both explicit tool-based retrieval and automatic context injection via a ContextProvider.
Changes:
- Introduces
BedrockKnowledgeBaseToolwithasync run()+get_tool_definition()for on-demand KB retrieval. - Introduces
BedrockKnowledgeBaseProvider(ContextProvider.before_run) to inject KB context automatically per agent invocation. - Adds a design/usage doc plus unit tests for the new tool/provider.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 17 comments.
| File | Description |
|---|---|
| python/packages/tools/agent_framework_tools/bedrock_knowledge_base.py | Adds the Bedrock Knowledge Base tool implementation and tool definition metadata. |
| python/packages/tools/agent_framework_tools/bedrock_knowledge_base_provider.py | Adds a context provider that retrieves and injects KB passages before runs. |
| python/packages/tools/tests/test_bedrock_knowledge_base.py | Adds unit tests covering tool retrieval, default config behavior, and provider retrieval. |
| python/packages/tools/agent_framework_tools/BEDROCK_MANAGED_KB.md | Adds documentation describing usage, configuration, and permissions. |
Comment on lines
+8
to
+16
| from agent_framework_tools.bedrock_kb import BedrockKnowledgeBaseTool | ||
|
|
||
| tool = BedrockKnowledgeBaseTool( | ||
| knowledge_base_id="YOUR_KB_ID", | ||
| region="us-east-1", | ||
| ) | ||
| results = tool.invoke({"query": "What are the compliance requirements?"}) | ||
| for result in results: | ||
| print(result["content"], result["score"]) |
| | AWS_REGION | AWS region for the KB | us-east-1 | | ||
| | AWS_PROFILE | AWS credentials profile | None | | ||
| | USE_AGENTIC_RETRIEVAL | Enable agentic retrieval | true | | ||
| | MAX_RESULTS | Maximum retrieval results | 5 | |
| @@ -0,0 +1,205 @@ | |||
| # Copyright (c) Amazon.com, Inc. All rights reserved. | |||
| @@ -0,0 +1,184 @@ | |||
| """Amazon Bedrock Knowledge Base retrieval tool for Microsoft Agent Framework. | |||
Comment on lines
+78
to
+81
| raise ImportError( | ||
| "boto3 is required for Bedrock Knowledge Base tool. " | ||
| "Install with: pip install boto3>=1.41.0" | ||
| ) |
Comment on lines
+194
to
+196
| except Exception as e: | ||
| logger.error(f"Error retrieving from Bedrock KB: {e}") | ||
| return [] |
Comment on lines
+2
to
+3
| from unittest.mock import MagicMock, patch | ||
| import pytest |
Comment on lines
+7
to
+21
| @patch("boto3.client") | ||
| def test_run_returns_results(self, mock_boto3): | ||
| from agent_framework_tools.bedrock_knowledge_base import BedrockKnowledgeBaseTool | ||
| mock_client = MagicMock() | ||
| mock_client.retrieve.return_value = { | ||
| "retrievalResults": [ | ||
| {"content": {"text": "Doc"}, "location": {"s3Location": {"uri": "s3://b/d"}}, "score": 0.9}, | ||
| ] | ||
| } | ||
| mock_boto3.return_value = mock_client | ||
| tool = BedrockKnowledgeBaseTool(knowledge_base_id="TEST123") | ||
| import asyncio | ||
| results = asyncio.run(tool.run(query="test")) | ||
| assert len(results) == 1 | ||
| assert results[0]["content"] == "Doc" |
Comment on lines
+23
to
+33
| @patch("boto3.client") | ||
| def test_managed_config_default(self, mock_boto3): | ||
| from agent_framework_tools.bedrock_knowledge_base import BedrockKnowledgeBaseTool | ||
| mock_client = MagicMock() | ||
| mock_client.retrieve.return_value = {"retrievalResults": []} | ||
| mock_boto3.return_value = mock_client | ||
| tool = BedrockKnowledgeBaseTool(knowledge_base_id="TEST123") | ||
| import asyncio | ||
| asyncio.run(tool.run(query="test")) | ||
| call_kwargs = mock_client.retrieve.call_args.kwargs | ||
| assert "managedSearchConfiguration" in call_kwargs["retrievalConfiguration"] |
Comment on lines
+54
to
+68
| @patch("boto3.client") | ||
| def test_retrieve_returns_passages(self, mock_boto3): | ||
| from agent_framework_tools.bedrock_knowledge_base_provider import BedrockKnowledgeBaseProvider | ||
| mock_client = MagicMock() | ||
| mock_client.retrieve.return_value = { | ||
| "retrievalResults": [ | ||
| {"content": {"text": "Context doc"}, "location": {"s3Location": {"uri": "s3://b/c"}}, "score": 0.9}, | ||
| ] | ||
| } | ||
| mock_boto3.return_value = mock_client | ||
| provider = BedrockKnowledgeBaseProvider(knowledge_base_id="TEST123") | ||
| import asyncio | ||
| passages = asyncio.run(provider._retrieve("test query")) | ||
| assert len(passages) == 1 | ||
| assert passages[0]["content"] == "Context doc" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
feat: add Amazon Bedrock Knowledge Base tool and context provider
Motivation & Context
Enables Agent Framework agents to retrieve context from Amazon Bedrock Knowledge Bases. This adds RAG capabilities using AWS's managed infrastructure without requiring users to manage their own vector stores or embedding pipelines.
Description & Review Guide
What are the major changes?
BedrockKnowledgeBaseTool— standalone tool withasync run()+get_tool_definition()for agents to call on-demandBedrockKnowledgeBaseProvider—ContextProvidersubclass withbefore_run()for automatic context injection on every agent invocationmanagedSearchConfiguration) and agentic retrieval with fallbackWhat is the impact of these changes?
python/packages/tools/, no existing code modifiedWhat do you want reviewers to focus on?
ContextProviderintegration withbefore_run()Related Issue
N/A — new feature adding Amazon Bedrock Knowledge Base integration.
Contribution Checklist
Testing
Tool.run()returned 5 results from live managed KB in us-west-2Provider._retrieve()returned 5 passages_format_context()produced 7399 chars of formatted contextboto3 >= 1.43